#include "StdAfx.h"
#include "AdvertBlocker.h"
#include "ElementHidingCollection.h"
#include "FilterListRegEntry.h"
#include <algorithm>
#include <iostream>
#include <fstream>
#include <boost/algorithm/string.hpp>
#include "wininet.h"
#include "RegistryManager.h"
#include "BufferForResource.h"
#include "time.h"
using namespace std;

CAdvertBlocker::CAdvertBlocker(void)
{
}

CAdvertBlocker::~CAdvertBlocker(void)
{
}

void CAdvertBlocker::LoadResources()
{
	LoadFilterManagerFromAppDataFolder();	
}

void CAdvertBlocker::LoadFilterManagerFromAppDataFolder()
{
	CRegistryManager regMgr;
	vector<CFilterListRegEntry> filterListRegEntries = regMgr.GetFilterListsRegEntries();

	for(auto it = filterListRegEntries.cbegin(); it != filterListRegEntries.cend(); ++it)
	{
		auto& filterListPath = it->Path;

		ifstream file;	
		file.open(filterListPath);
	
		auto filterList = make_shared<CFilterList>();

		if(file.is_open())
		{			
			boost::archive::text_iarchive ia(file);
			ia >> *filterList;			
		}
		file.close();

		_filterManager.AddFilterList(filterList);
	}
}

void CAdvertBlocker::Initialize(CComPtr<IWebBrowser2> webBrowser)
{
	CComPtr<IDispatch> dispDoc;
    HRESULT hr = webBrowser->get_Document(&dispDoc);
    if (SUCCEEDED(hr))
	{
		CComQIPtr<IHTMLDocument2> htmlDoc2 = dispDoc;
		if(htmlDoc2)
		{
			CreateClassElementsMap(htmlDoc2);
		}
	}
}

void CAdvertBlocker::CleanUpAfterInitialization()
{
	_elementsOfClass.clear();
}

void CAdvertBlocker::CreateClassElementsMap(CComPtr<IHTMLDocument2> document)
{
	CComPtr<IHTMLElementCollection> elements;
	document->get_all(&elements);
	if(elements)
	{
		long length;
		elements->get_length(&length);
		for(int i = 0; i < length; i++)
		{
			CComVariant itemIndex(i);
			CComPtr<IDispatch> dispElement;
			HRESULT hr = elements->item(itemIndex, itemIndex, &dispElement);
			if(SUCCEEDED(hr) && (dispElement))
			{
				CComQIPtr<IHTMLElement> htmlElement = dispElement;
				if(htmlElement)
				{
					CComBSTR classNameBstr;
					htmlElement->get_className(&classNameBstr);
					if(classNameBstr != NULL)
					{
						hr = classNameBstr.ToLower();
						if(SUCCEEDED(hr))
						{
							wstring className(classNameBstr);
							_elementsOfClass[className].push_back(htmlElement);
						}
					}
				}
			}
		}
	}
}

void CAdvertBlocker::CreateFilterManagerAndSerializeToFile(string pathToBlockingRules, string pathToSerializedFileMgr)
{
	////creating FilterManager from file
	//CFilterManager fm;	
	//fm.InitializeFromAdblockPlusFiltersFile(pathToBlockingRules);

	////serializing
	//ofstream fileStream(pathToSerializedFileMgr);
	//boost::archive::text_oarchive outputArchive(fileStream);
	//outputArchive << fm;
}

void CAdvertBlocker::RemoveAdverts(CComPtr<IWebBrowser2> webBrowser)
{
	CComPtr<IDispatch> dispDoc;
    HRESULT hr = webBrowser->get_Document(&dispDoc);
	if(SUCCEEDED(hr) && dispDoc)
	{
		CComQIPtr<IHTMLDocument3> doc3 = dispDoc;
		if(doc3)
		{
			RemoveAdverts(doc3);
		}
	}
}

void CAdvertBlocker::RemoveAdverts(CComPtr<IHTMLDocument3> document3)
{	
	RemoveAdvertsBasedOnElementHidingRules(document3);

	RemoveObjectsAdverts(document3); //flash objects 

	RemoveAdvertsWithSrcAttribute(document3, CComBSTR(L"img"));
	RemoveAdvertsWithSrcAttribute(document3, CComBSTR(L"iframe"));
	RemoveAdvertsWithSrcAttribute(document3, CComBSTR(L"embed"));
	RemoveAdvertsWithSrcAttribute(document3, CComBSTR(L"script"));	
}

void CAdvertBlocker::RemoveAdvertsBasedOnElementHidingRules(CComPtr<IHTMLDocument3> document3)
{	
	CComQIPtr<IHTMLDocument6> document6 = document3;
	if(document6)
	{
		vector<shared_ptr<CElementHidingCollection>> allElementHidingRules = _filterManager.GetElementHidingCollections();

		for(auto elemsInList = allElementHidingRules.begin(); elemsInList != allElementHidingRules.end(); ++elemsInList)
		{
			shared_ptr<CElementHidingCollection> elems = (*elemsInList);
			
			
			for(auto it = elems->GetAdvertIdElements().cbegin();
			it != elems->GetAdvertIdElements().cend(); ++it)
			{
				RemoveHtmlElementWithId(document6, *it);		
			}

			for(auto it = elems->GetAdvertClassElements().cbegin();
				it != elems->GetAdvertClassElements().cend(); ++it)
			{
				RemoveHtmlElementWithClass(document6,*it);
			}	
		}		
	}
}

void CAdvertBlocker::RemoveHtmlElementWithId(CComPtr<IHTMLDocument6> document, const wstring& id)
{
	CComPtr<IHTMLElement2> element;
	CComBSTR idBstr(id.c_str());

	HRESULT hr = document->getElementById(idBstr, &element);
	if(SUCCEEDED(hr) && element != NULL)
	{
		CComQIPtr<IHTMLDOMNode> node = element;
		RemoveHtmlDOMNode(node);
	}	
}

void CAdvertBlocker::RemoveHtmlElementWithClass(CComPtr<IHTMLDocument6> document, const wstring& cl)
{

	if(_elementsOfClass.count(cl))	//if there is any alement of this class
	{
		auto& elements = _elementsOfClass[cl];	//gets vector of elements
		for(auto iter = elements.begin(); iter != elements.end(); ++iter)
		{
			RemoveHtmlElement(*iter);
		}
		elements.clear();
	}
}

void CAdvertBlocker::RemoveObjectsAdverts(CComPtr<IHTMLDocument3> document3)
{
	CComPtr<IHTMLElementCollection> elements;
	HRESULT hr = document3->getElementsByTagName(CComBSTR(L"object"), &elements);
	
	if(SUCCEEDED(hr) && elements != NULL)
	{
		RemoveHtmlElements(elements, &CAdvertBlocker::IsObjectAdvert);
	}
}

bool CAdvertBlocker::IsObjectAdvert(CComPtr<IHTMLElement> element)
{
	wstring objectSrc = GetSrcFromObjectElement(element);
	boost::to_lower(objectSrc);

	return _filterManager.MatchAdvertFilter(objectSrc);
}

wstring CAdvertBlocker::GetSrcFromObjectElement(CComPtr<IHTMLElement> object)
{	
#ifdef DEBUG
	CComBSTR bstr;
	object->get_innerHTML(&bstr);
#endif

	HRESULT hr = S_OK;

	CComQIPtr<IHTMLDOMNode> node = object;
	CComPtr<IHTMLDOMNode> firstChild;

	hr = node->get_firstChild(&firstChild);

	if( SUCCEEDED(hr) && (firstChild) )
	{
		if(IsParamTagWithNameMovieAttribute(firstChild))
		{
			wstring value = GetValueAttribute(firstChild);
			return value;
		}

		CComPtr<IHTMLDOMNode> sibling;
		hr = firstChild->get_nextSibling(&sibling);
		while(SUCCEEDED(hr) && sibling)
		{
			if(IsParamTagWithNameMovieAttribute(sibling))
			{
				wstring value = GetValueAttribute(sibling);
				return value;
			}
			CComPtr<IHTMLDOMNode> tmpNode;
			hr = sibling->get_nextSibling(&tmpNode);
			sibling = tmpNode;
		}
	}

	return wstring();
}

bool CAdvertBlocker::IsParamTagWithNameMovieAttribute(CComPtr<IHTMLDOMNode> node)
{
	// checking if this node is like this:
	//<param name="movie" value="flash_address.swf"/>

	if(IsParamTag(node) && HasNameAttributeWithMovieValue(node))
		return true;
	else
		return false;
}

bool CAdvertBlocker::IsParamTag(CComPtr<IHTMLDOMNode> node)
{
	CComQIPtr<IHTMLElement> element = node;
	if(!element)
	{
		return false;
	}
	else
	{
		CComBSTR bstr;
		element->get_tagName(&bstr);
		if(bstr != NULL)
		{
			HRESULT hr = bstr.ToLower();
			if(SUCCEEDED(hr))
			{
				if(wstring(bstr) == L"param")
				{
					return true;
				}
			}
		}
	}
	return false;
}

bool CAdvertBlocker::HasNameAttributeWithMovieValue(CComPtr<IHTMLDOMNode> node)
{
	CComPtr<IDispatch> dispAttr;
	node->get_attributes(&dispAttr);

	if(dispAttr)
	{		
		CComQIPtr<IHTMLAttributeCollection3> attributes = dispAttr;

		if(attributes)
		{
			CComBSTR nameBstr(L"name");
			CComPtr<IHTMLDOMAttribute> nameAttr;

			HRESULT hr = attributes->getNamedItem(nameBstr, &nameAttr);			
			if( SUCCEEDED(hr) && (nameAttr) )
			{
				CComVariant valueVariant;
				hr = nameAttr->get_nodeValue(&valueVariant);
				if(SUCCEEDED(hr) && (valueVariant.vt != VT_NULL) && (valueVariant.bstrVal != NULL))
				{
					CComBSTR nameBstr(valueVariant.bstrVal);
					hr = nameBstr.ToLower();
					if(SUCCEEDED(hr))
					{
						wstring nameValue(nameBstr);
						
						if(nameValue == L"movie")
						{
							return true;
						}
					}
				}
			}
		}
	}
	return false;
}

wstring CAdvertBlocker::GetValueAttribute(CComPtr<IHTMLDOMNode> param)
{
	CComPtr<IDispatch> dispAttr;
	param->get_attributes(&dispAttr);

	if(dispAttr)
	{		
		CComQIPtr<IHTMLAttributeCollection3> attributes = dispAttr;

		if(attributes)
		{
			CComBSTR valueBstr(L"value");
			CComPtr<IHTMLDOMAttribute> valueAttr;
							
			HRESULT hr = attributes->getNamedItem(valueBstr, &valueAttr);
			if(SUCCEEDED(hr))
			{
				CComVariant valueVariant;
				valueAttr->get_nodeValue(&valueVariant);
				if(valueVariant.vt != VT_NULL)
				{
					wstring value(CComBSTR(valueVariant.bstrVal));
					return value;
				}			
			}
		}
	}

	return wstring(); //returning empty string if something goes wrong
}

void CAdvertBlocker::RemoveAdvertsWithSrcAttribute(CComPtr<IHTMLDocument3> document3, CComBSTR elementName)
{
	CComPtr<IHTMLElementCollection> elements;
	HRESULT hr = document3->getElementsByTagName(elementName, &elements);
	if(SUCCEEDED(hr) && elements != NULL)
	{
		RemoveHtmlElements(elements, &CAdvertBlocker::IsSrcAttributeAnAdvert);
	}
}

bool CAdvertBlocker::IsSrcAttributeAnAdvert(CComPtr<IHTMLElement> element)
{	
	CComQIPtr<IHTMLElement5> elementWithSrcAttribute = element;
	if(elementWithSrcAttribute)
	{
		CComBSTR srcBstr(L"src");
		CComVariant srcVariant;

		HRESULT hr = elementWithSrcAttribute->getAttribute(srcBstr, &srcVariant);

		//null BSTR are considered to be empty strings
		if( SUCCEEDED(hr) && (srcVariant.vt != VT_NULL) && (srcVariant.bstrVal != NULL) ) 
		{
			CComBSTR srcBstr(srcVariant.bstrVal);
			wstring src(srcBstr);
			boost::to_lower(src);
			return _filterManager.MatchAdvertFilter(src);
		}
	}
	return false;
}

void CAdvertBlocker::RemoveHtmlElements(CComPtr<IHTMLElementCollection> elements,  FilterFunction isAdvert)
{	
	long length;
	elements->get_length(&length);
	for(int i = 0; i < length; i++)
	{
		CComVariant itemIndex(i);
		CComVariant empty;
		CComPtr<IDispatch> dispElement;

		HRESULT hr = elements->item(itemIndex, itemIndex, &dispElement);
		if(SUCCEEDED(hr) && dispElement != NULL)
		{
			CComQIPtr<IHTMLElement> htmlElement = dispElement;
			if(htmlElement)
			{
				if( (isAdvert == NULL) || ((this->*isAdvert)(htmlElement)) )
				{
					CComQIPtr<IHTMLDOMNode> htmlDomNode = htmlElement;
					if(htmlDomNode)
					{
						RemoveHtmlDOMNode(htmlDomNode);
					}
				}
			}
		}
	}
}

void CAdvertBlocker::RemoveHtmlDOMNode(CComPtr<IHTMLDOMNode> node)
{
	CComPtr<IHTMLDOMNode> ptr;
	node->removeNode(VARIANT_TRUE, &ptr);
}

void CAdvertBlocker::RemoveHtmlElement(CComPtr<IHTMLElement> element)
{
	CComQIPtr<IHTMLDOMNode> node = element;
	RemoveHtmlDOMNode(node);
}

